#!/usr/bin/env python3
# set coding: utf-8
"""
@Project : litewaf2
@File :    views.py
@Time :    2023/2/21 10:54
@Author :  richard zhu
@Email :   gaotao.zhugt@dtzhejiang.com
@purpose : 
"""

from django.views.decorators.csrf import csrf_exempt
from django.http import JsonResponse, HttpResponse
from django.db.models import Q
from django.shortcuts import render
from django.views.generic import View
from captcha.models import CaptchaStore
from captcha.helpers import captcha_image_url
from api.models import WafUser, Node, GlobalSetting, WafDomain, WafRule, FlowRule, RuleGroup
from utils import get_real_ip
import traceback
import logging
import json
import datetime
import pytz
import redis


logger = logging.getLogger('litewaf')

class MyException(Exception):
    pass


def __get_global_settings():
    global_rule_data = GlobalSetting.objects.filter(status=1).last()
    if global_rule_data:
        return global_rule_data
    else:
        raise MyException("not found global setting")

def require_login():
    """
    登陆校验
    """
    def _deco(func):
        def __deco(request, *args, **kwargs):
            data_result = {}
            try:
                assert request.method == 'POST'
                json_data = json.loads(request.body)
                api_key = json_data['api_key']
                api_password = json_data['api_password']
            except AssertionError:
                data_result['resutl'] = False
                data_result['message'] = "ONLY support method POST"
                return JsonResponse(data_result, status=403)
            except Exception as e:
                data_result['resutl'] = False
                data_result['message'] = "parameter error:" + str(e)
                return JsonResponse(data_result, status=400)

            try:
                user = WafUser.objects.get(Q(api_key=api_key)&Q(api_password=api_password)&Q(status=1))
            except Exception as e:
                logger.error("user api login failed, %s[%s]" % (api_key,api_password))
                data_result['result'] = False
                data_result['message'] = "api_key or api_password error"
                return JsonResponse(data_result, status=401)
            return func(request, user, json_data, *args, **kwargs)
        return __deco
    return _deco


# @csrf_exempt
@require_login()
def heartbeat(request, user, json_data):
    data_result = {}
    data_result['resutl'] = True
    data_result['message'] = "heartbeat update successful"
    try:
        node_uuid = json_data['node_uuid']
        node_hostname = json_data['node_hostname']
    except Exception as e:
        data_result['resutl'] = False
        data_result['message'] = "parameter error:" + str(e)
        return JsonResponse(data_result, status=400)

    try:
        Node.objects.update_or_create(uuid=node_uuid,
                                      defaults={'ip': get_real_ip(request),
                                                'gmt_modify': datetime.datetime.now(),
                                                'hostname': node_hostname, 'user_id': user})
    except Exception as e:
        data_result['resutl'] = False
        data_result['message'] = "parameter error:" + str(e)
        logger.error(user)
        return JsonResponse(data=data_result, status=500)

    return JsonResponse(data=data_result, status=200)


@require_login()
def config_global(request, user, json_data):
    data_result = {}

    global_rule_data = __get_global_settings()
    keeptime = (datetime.datetime.now() - datetime.timedelta(seconds=1800)).replace(tzinfo=pytz.timezone('UTC'))

    # 半小时内，客户端需同步
    if keeptime > global_rule_data.gmt_modify:
        data_result['configure_without_change'] =  False
        data_result['global_rule_data'] = global_rule_data.to_dict()
        data_result['result'] = True
    else:
        data_result['configure_without_change'] =  True
        data_result['global_rule_data'] = {}
        data_result['result'] = True

    return JsonResponse(data=data_result, status=200)


def m2m_nested(rules, rgs, deep_cnt):
    rgrgs = RuleGroup.objects.none()
    for rg in rgs:
        if rules.count() > 0:
            rules = rules | rg.rule.filter(status=True)
        else:
            rules = rg.rule.filter(status=True)

        rgrgs = rgrgs | rg.rg.filter(status=True)

    if rgrgs.count() == 0:
        return rules

    deep_cnt -= 1
    while deep_cnt < 0:
        raise RecursionError("递归深度过大，检查是否有循环")
    return m2m_nested(rules, rgrgs, deep_cnt)


@require_login()
def config_domain(request, user, json_data):
    data_result = {}
    keeptime = (datetime.datetime.now() - datetime.timedelta(seconds=1800)).replace(tzinfo=pytz.timezone('UTC'))

    waf_rule_kt = WafRule.objects.filter(gmt_modify__lt=keeptime).count()
    flow_rule_kt = FlowRule.objects.filter(gmt_modify__lt=keeptime).count()
    rg_kt = RuleGroup.objects.filter(gmt_modify__lt=keeptime).count()

    # 半小时内，Rule,RuleGroup没有更新则
    if waf_rule_kt == 0 and flow_rule_kt == 0  and rg_kt == 0:
        data_result['configure_without_change'] = True
    else:
        data_result['configure_without_change'] = False

    try:
        domain_dict = {}
        rules_data = {}

        for dom in WafDomain.objects.filter(status=True).filter(user_id=user):
            rule_dict = {}
            rgs = dom.rulegroup.filter(status=True)
            domain_with_port = "%s:%s" % (dom.domain, dom.port)
            domain_dict[domain_with_port] = {"rules": rule_dict}

            for rg in rgs:
                waf_rules = rg.waf_rule.filter(status=True).order_by('order')
                for _rule in waf_rules:
                    if rule_dict.get(_rule.rule_type):
                        rule_dict[_rule.rule_type].append(_rule.uuid)
                    else:
                        rule_dict[_rule.rule_type] = [_rule.uuid]

                    if rules_data.get(_rule.rule_type):
                        rules_data[_rule.rule_type].update({_rule.uuid: _rule.to_dict()})
                    else:
                        rules_data[_rule.rule_type] = {_rule.uuid: _rule.to_dict()}

                flow_rules = rg.flow_rule.filter(status=True).order_by('order')

                for _rule in flow_rules:
                    if rule_dict.get(_rule.rule_type):
                        rule_dict[_rule.rule_type].append(_rule.uuid)
                    else:
                        rule_dict[_rule.rule_type] = [_rule.uuid]
                    if rules_data.get(_rule.rule_type):
                        rules_data[_rule.rule_type].update({_rule.uuid: _rule.to_dict()})
                    else:
                        rules_data[_rule.rule_type] = {_rule.uuid: _rule.to_dict()}

        data_result['domain_rule_data'] = domain_dict
        data_result['rules_data'] = rules_data
        data_result['configure_without_change'] = False
        data_result['auto_update_period'] = __get_global_settings().auto_update_period
        data_result['result'] = True
        return JsonResponse(data=data_result, status=200)

    except Exception as e:
        data_result['resutl'] = False
        data_result['message'] = "parameter error:" + str(e)
        logger.error(user)
        return JsonResponse(data=data_result, status=500)


# 创建验证码
def captcha():
    hashkey = CaptchaStore.generate_key()   #验证码答案
    image_url = captcha_image_url(hashkey)  #验证码地址
    captcha = {'hashkey': hashkey, 'image_url': image_url}
    return captcha
# 刷新验证码
def refresh_captcha(request):
    return HttpResponse(json.dumps(captcha()), content_type='application/json')

# 验证验证码
def jarge_captcha(captchaStr, captchaHashkey):
    if captchaStr and captchaHashkey:
        try:
            # 获取根据hashkey获取数据库中的response值
            get_captcha = CaptchaStore.objects.get(hashkey=captchaHashkey)
            if get_captcha.response == captchaStr.lower():     # 如果验证码匹配
                return True
        except:
            return False
    else:
        return False


def bot_check_in_redis(rule_id):
    rule_type = "bot_check"
    redis_key = rule_type+":"+rule_id
    get_global_data = __get_global_settings()
    redis_config = get_global_data.redis
    redis_client = redis.StrictRedis(host=redis_config['host'], port=redis_config['port'], password=redis_config['password'])
    if redis_client.ping():
        return redis_client.set(redis_key, 1)
    return False


# 验证码页面
class CaptchaView(View):
    def get(self, request):
        hashkey = CaptchaStore.generate_key()  # 验证码答案
        image_url = captcha_image_url(hashkey)  # 验证码地址
        rule_id = request.GET.get("rule_id")        # 获取ruleid用于匹配redis规则
        captcha = {'hashkey': hashkey, 'image_url': image_url, 'rule_id': rule_id}
        return render(request, "captcha.html", locals())

    def post(self,request):
        capt=request.POST.get("captcha",None)         #用户提交的验证码
        key=request.POST.get("hashkey",None)          #验证码答案
        rule_id=request.POST.get("rule_id", None)
        if jarge_captcha(capt,key):
            ret = bot_check_in_redis(rule_id)
            logger.info("bot_check_in_redis",ret)
            if ret:
                logger.info("bot_check_in_redis set ok")
            else:
                logger.error("bot_check_in_redis set error")

            return  HttpResponse("验证码正确")
        else:
            return HttpResponse("验证码错误")
